#-*- coding:utf-8 -*-
import sys
reload(sys)
sys.setdefaultencoding('utf-8')
import volatility.conf as conf
import volatility.registry as registry
import volatility.plugins.vadinfo as vadinfo
import volatility.plugins.dlldump as dlldump
import volatility.commands as commands
import volatility.addrspace as addrspace
import volatility.plugins.envars as envars
import json
import pefile

class VadAnalyse:
	def GetVadinfo(self,src_path,dst_path):
		self.src_path = src_path
		self.dst_path = dst_path
		registry.PluginImporter()
		config = conf.ConfObject()
		registry.register_global_options(config, commands.Command)
		registry.register_global_options(config, addrspace.BaseAddressSpace)
		config.parse_options()
		config.PROFILE="WinXPSP2x86"
		config.LOCATION = "file://{}".format(src_path)
		config.OUTPUT = "json"
		dst_path = "{}vadinfo.json".format(dst_path)
		config.OUTPUT_FILE = dst_path
		info = vadinfo.VADInfo(config)
		info.execute()
		with open(dst_path,"r") as f:
			data = json.load(f)
			f.close()
		data = data["rows"]
		head = ['pid', 'vad', 'start', 'end', 'VadTag', 'flags', 'protection', 'VadType', 'ControlArea', 'segment', \
               'NumberOfSectionReferences', 'NumberOfPfnReferences', 'NumberOfMappedViews', 'NumberOfUserReferences',\
               'ControlFlags', 'FileObject', 'FileName', 'FirstprototypePTE', 'LastcontiguousPTE', 'Flags2']
		result = dict()
		for vad_item in data:
			temp = dict(zip(head,vad_item))
			if not result.has_key(temp["pid"]):
				result[temp["pid"]] = list()
				result[temp["pid"]].append(0)
			if temp["FileName"][-4:] == ".exe":
				result[temp["pid"]][0] = temp["FileObject"]
			result[temp["pid"]].append(temp)
		with open(dst_path,"w") as f:
			f.write(json.dumps(result))
			f.close()

	def VadDump(self,src_path,dst_path):
		registry.PluginImporter()
		config = conf.ConfObject()
		registry.register_global_options(config, commands.Command)
		registry.register_global_options(config, addrspace.BaseAddressSpace)
		config.parse_options()
		config.PROFILE="WinXPSP2x86"
		config.LOCATION = "file://{}".format(src_path)
		config.DUMP_DIR = dst_path
		dump = vadinfo.VADDump(config)
		dump.execute()

	def ParsePE(self,src_path,dst_path):
		pe = pefile.PE(src_path)
		info_dict = pe.dump_dict("all")
		with open(dst_path,"w") as fp:
			fp.write(json.dumps(info_dict))
